import numpy as np
import cv2 as cv
import math
import glob
import os
from math import *

# 标定板 9*6 25mm
cols = 10 - 1
rows = 7 - 1
distance = 25


# 5张标定板固定的图片计算的相机内参矩阵和畸变系数
cam_mtx = np.array([[2.41470924e+03, 0.00000000e+00, 1.30239360e+03], 
[0.00000000e+00, 2.41454060e+03, 7.79527310e+02], 
[0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], dtype=np.float64)

Distortion = np.array([0.1546888, -0.76913489, 0.00362884, 0.00357295, 0.84778932], dtype=np.float64)

# # 20张标定板固定的图片计算的相机内参矩阵和畸变系数
# cam_mtx = np.array([[2.39878870e+03, 0.00000000e+00, 1.32335648e+03],
# [0.00000000e+00, 2.39881495e+03, 7.52183972e+02], 
# [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], dtype=np.float64)

# Distortion = np.array([1.36807510e-01, -7.10167547e-01, -8.04371737e-04, 3.04634395e-03, 9.71742348e-01], dtype=np.float64)


# termination criteria
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
# 标定板世界坐标
objp = np.zeros((rows*cols,3), np.float32)
for m in range(rows):
    for n in range(cols):
        objp[m*cols + n] = [n*distance, m*distance, 0]

# Arrays to store object points and image points from all the images.
objpoints = [] # 3d point in real world space
imgpoints = [] # 2d points in image plane.
images = glob.glob('calibration/cam_clib/pictures/*.jpg')
cv.namedWindow('img', 0)
for fname in images:
    img = cv.imread(fname)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    
    # Find the chess board corners
    ret, corners = cv.findChessboardCorners(gray, (cols, rows), None)
    # If found, add object points, image points (after refining them)
    if ret == True:
        objpoints.append(objp)
        corners2 = cv.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        # Draw and display the corners
        cv.drawChessboardCorners(img, (cols, rows), corners2, ret)
        cv.namedWindow('img')
        cv.resizeWindow('img', 600, 600)
        cv.moveWindow('img',300,300)
        cv.imshow('img', img)
        cv.waitKey(200)
        # input("请输入任意字符继续：")
# 相机校准 返回参数：相机矩阵、畸变系数、旋转向量和平移向量
# dist-畸变系数   distortion cofficients = (k_1,k_2,p_1,p_2,k_3)
'''
得出的rvec怎么是3*1的矩阵?
调用solvePnPRansac()函数得到的rvec是一个旋转矢量,需要使用cv2.Rodrigues(src,dst,jacobian=None)
src:输入的矩阵可以是(3*1或1*3)也可以是3*3
dst:输出的矩阵,对应输入3*3,或者(3*1或1*3)
jacobian:也是一个输出，该输出表明了输入矩阵和输出矩阵的的雅可比
转换后的矩阵就和前面的旋转矩阵对应了
'''
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
# 与calibrateCamera()运行的结果相同
# for i in range(20):
#     print('第'+str(i)+'张')
#     retval,rvec,tvec = cv.solvePnP(objpoints[i],imgpoints[i], mtx, dist)
#     print('旋转向量slovePnp:\n',rvec)
#     print('平移向量slovePnp：\n',tvec)
#     print('标定得到的旋转向量:\n',rvecs[i])
#     print('标定得到的平移向量:\n',tvecs[i])
#     print('')
    # 反投影误差
total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv.norm(imgpoints[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
    total_error += error
print("total error: ", total_error / len(objpoints))

print('内参矩阵:\n',mtx)
print('畸变矩阵:\n',dist)

# 测试代码

img = cv.imread('calibration/cam_clib/pictures/0.jpg')
h, w = img.shape[:2]
newcameramtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))
# undistort
dst = cv.undistort(img, mtx, dist, None, newcameramtx)
# crop the image
x, y, w, h = roi
dst = dst[y:y+h, x:x+w]
cv.imwrite('calibration/cam_clib/calib_results/calibresult.png', dst)
cv.destroyAllWindows()

